package com.wfh.frp;

import com.wfh.frp.server.MainServer;
import com.wfh.frp.server.ProxyClientServer;
import com.wfh.frp.server.VisitorServer;
import lombok.extern.slf4j.Slf4j;

import java.util.concurrent.CountDownLatch;

/**
 * Internal Network Penetration Tool
 * Supports server, client, and visitor modes with configurable ports
 *
 * Usage:
 *   java -jar myfrp.jar server [port]                    # Start as server on specified port (default: 9090)
 *   java -jar myfrp.jar client <serverHost> <serverPort> <targetHost> <targetPort> [localPort]  # Start as client
 *   java -jar myfrp.jar visitor <port> <serverHost> <serverPort>  # Start as visitor on specified port
 *
 * @author fenghuanwang
 */
@Slf4j
public class MyFrp {
    private static final int DEFAULT_SERVER_PORT = 9090;
    private static final int DEFAULT_VISITOR_PORT = 8081;
    private static final String LOCALHOST = "127.0.0.1";

    public static void main(String[] args) throws InterruptedException {
        if (args.length == 0) {
            printUsageAndExit();
        }

        String mode = args[0].toLowerCase();

        switch (mode) {
            case "server":
                handleServerMode(args);
                break;
            case "client":
                handleClientMode(args);
                break;
            case "visitor":
                handleVisitorMode(args);
                break;
            default:
                log.error("Unknown mode: {}. Use 'server', 'client', or 'visitor'", mode);
                printUsageAndExit();
        }

        log.info("服务已经启动。。。");
        Runtime.getRuntime().addShutdownHook(new Thread(() -> log.info("服务正在关闭。。。")));
        new CountDownLatch(1).await();
    }

    private static void handleServerMode(String[] args) {
        int port = DEFAULT_SERVER_PORT;
        if (args.length > 1) {
            try {
                port = Integer.parseInt(args[1]);
            } catch (NumberFormatException e) {
                log.error("Invalid port number: {}. Using default port {}", args[1], port);
            }
        }

        // Use MainServer's start method with port parameter
        int finalPort = port;
        Thread.startVirtualThread(() -> MainServer.start(finalPort));
    }

    private static void handleClientMode(String[] args) {
        if (args.length < 5) {
            log.error("Client mode requires at least 4 arguments: <serverHost> <serverPort> <targetHost> <targetPort>");
            printUsageAndExit();
        }

        String serverHost = args[1];
        int serverPort;
        String targetHost = args[3];
        int targetPort;

        try {
            serverPort = Integer.parseInt(args[2]);
            targetPort = Integer.parseInt(args[4]);
        } catch (NumberFormatException e) {
            log.error("Invalid port number. Please provide valid integers for ports.");
            System.exit(1);
            return; // This won't be reached but keeps compiler happy
        }

        // Optional local port for the client
        int localPort = args.length > 5 ? Integer.parseInt(args[5]) : 9090;

        Thread.startVirtualThread(() -> ProxyClientServer.start(serverHost, serverPort, targetHost, targetPort));
    }

    private static void handleVisitorMode(String[] args) {
        if (args.length < 4) {
            log.error("Visitor mode requires 3 arguments: <port> <serverHost> <serverPort>");
            printUsageAndExit();
        }

        int port;
        String serverHost = args[2];
        int serverPort;

        try {
            port = Integer.parseInt(args[1]);
            serverPort = Integer.parseInt(args[3]);
        } catch (NumberFormatException e) {
            log.error("Invalid port number. Please provide valid integers for ports.");
            System.exit(1);
            return; // This won't be reached but keeps compiler happy
        }

        Thread.startVirtualThread(() -> VisitorServer.start(port, serverHost, serverPort));
    }

    private static void printUsageAndExit() {
        log.info("\nUsage:\n" +
                "  java -jar myfrp.jar server [port]                    # Start as server on specified port (default: 9090)\n" +
                "  java -jar myfrp.jar client <serverHost> <serverPort> <targetHost> <targetPort> [localPort]  # Start as client\n" +
                "  java -jar myfrp.jar visitor <port> <serverHost> <serverPort>  # Start as visitor on specified port\n");
        System.exit(1);
    }
}